#!/usr/bin/env python3
"""
受试者目录结构创建工具
自动为每个受试者创建标准的文件夹结构
"""

import os
import argparse
from pathlib import Path
import logging
from typing import List

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class SubjectDirectoryCreator:
    """受试者目录结构创建器"""
    
    def __init__(self, base_dir: str):
        self.base_dir = Path(base_dir)
        self.logger = logging.getLogger(__name__)
        
        # 定义主目录结构
        self.main_dirs = {
            'prediction': self.base_dir / 'prediction',
            'segmentation': self.base_dir / 'segmentation',
            'reference': self.base_dir / 'reference',
            'resliced': self.base_dir / 'resliced',
            'metric_results': self.base_dir / 'metric_results'
        }
    
    def create_subject_directories(self, subject_ids: List[str], dry_run: bool = False) -> bool:
        """
        为指定的受试者自动创建标准文件夹结构
        
        目录结构:
        ├── prediction/
        │   ├── subject1/
        │   │   ├── processed/
        │   │   ├── proposed/
        │   │   ├── source/
        │   │   └── contrast_methods/
        │   └── subject2/...
        ├── segmentation/
        │   ├── subject1/
        │   └── subject2/...
        ├── reference/
        │   ├── subject1/
        │   └── subject2/
        ├── resliced/
        │   ├── subject1/
        │   │   ├── images/
        │   │   └── masks/
        │   └── subject2/...
        ├── metric_results/
        
        Args:
            subject_ids: 受试者ID列表
            dry_run: 是否为模拟运行，True时只显示将要创建的目录
            
        Returns:
            创建是否成功
        """
        if not subject_ids:
            self.logger.warning("未提供受试者ID列表")
            return False
        
        # 定义每个主目录下需要创建的子目录结构
        directory_structure = {
            'prediction': ['processed', 'proposed', 'source', 'contrast_methods'],
            'segmentation': [],  # 只创建subject目录
            'reference': [],     # 只创建subject目录
            'resliced': ['images', 'masks']
        }
        
        self.logger.info(f"{'='*50}")
        self.logger.info(f"受试者目录结构创建工具")
        self.logger.info(f"{'='*50}")
        self.logger.info(f"基础目录: {self.base_dir}")
        self.logger.info(f"受试者列表: {', '.join(subject_ids)}")
        
        if dry_run:
            self.logger.info("\n[模拟运行] 将要创建以下目录结构:")
            self._show_directory_preview(subject_ids, directory_structure)
            return True
        
        success = True
        created_dirs = []
        
        # 首先确保主目录存在
        for main_dir_name, main_dir_path in self.main_dirs.items():
            try:
                main_dir_path.mkdir(parents=True, exist_ok=True)
                if main_dir_name != 'metric_results':  # metric_results不需要subject子目录
                    self.logger.info(f"✓ 确保主目录存在: {main_dir_path}")
            except Exception as e:
                self.logger.error(f"✗ 创建主目录失败 {main_dir_path}: {e}")
                success = False
        
        # 为每个受试者创建目录结构
        for subject_id in subject_ids:
            self.logger.info(f"\n--- 创建受试者目录: {subject_id} ---")
            
            for main_dir, subdirs in directory_structure.items():
                # 创建主受试者目录
                subject_dir = self.main_dirs[main_dir] / subject_id
                
                try:
                    subject_dir.mkdir(parents=True, exist_ok=True)
                    created_dirs.append(str(subject_dir))
                    self.logger.info(f"  ✓ {subject_dir}")
                except Exception as e:
                    self.logger.error(f"  ✗ 创建目录失败 {subject_dir}: {e}")
                    success = False
                    continue
                
                # 创建子目录
                for subdir in subdirs:
                    sub_path = subject_dir / subdir
                    
                    try:
                        sub_path.mkdir(parents=True, exist_ok=True)
                        created_dirs.append(str(sub_path))
                        self.logger.info(f"    ✓ {sub_path}")
                    except Exception as e:
                        self.logger.error(f"    ✗ 创建子目录失败 {sub_path}: {e}")
                        success = False
        
        # 确保metric_results目录存在
        try:
            self.main_dirs['metric_results'].mkdir(parents=True, exist_ok=True)
            self.logger.info(f"✓ 结果目录: {self.main_dirs['metric_results']}")
        except Exception as e:
            self.logger.error(f"✗ 创建结果目录失败: {e}")
            success = False
        
        self.logger.info(f"\n{'='*50}")
        if success:
            self.logger.info(f"✓ 目录创建完成! 总共创建了 {len(created_dirs)} 个目录")
        else:
            self.logger.warning("⚠ 部分目录创建失败，请检查权限和路径")
        self.logger.info(f"{'='*50}")
        
        return success
    
    def _show_directory_preview(self, subject_ids: List[str], directory_structure: dict):
        """显示将要创建的目录结构预览"""
        self.logger.info(f"\n项目根目录: {self.base_dir}/")
        
        for main_dir, subdirs in directory_structure.items():
            self.logger.info(f"├── {main_dir}/")
            for i, subject_id in enumerate(subject_ids):
                is_last_subject = (i == len(subject_ids) - 1)
                subject_prefix = "└──" if is_last_subject else "├──"
                self.logger.info(f"│   {subject_prefix} {subject_id}/")
                
                for j, subdir in enumerate(subdirs):
                    is_last_subdir = (j == len(subdirs) - 1)
                    subdir_prefix = "└──" if is_last_subdir else "├──"
                    indent = "    " if is_last_subject else "│       "
                    self.logger.info(f"│   {indent}{subdir_prefix} {subdir}/")
        
        self.logger.info(f"└── metric_results/")


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="受试者目录结构创建工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 为单个受试者创建目录结构
  python create_subject_dirs.py --subjects CHEN-REN-GENG
  
  # 为多个受试者创建目录结构
  python create_subject_dirs.py --subjects CHEN-REN-GENG DJY BCR
  
  # 模拟运行，查看将要创建的目录结构
  python create_subject_dirs.py --subjects CHEN-REN-GENG --dry-run
  
  # 指定项目根目录
  python create_subject_dirs.py --base-dir /path/to/project --subjects CHEN-REN-GENG
        """
    )
    
    parser.add_argument(
        "--base-dir", 
        type=str, 
        default="..", 
        help="项目根目录路径 (默认: ..)"
    )
    parser.add_argument(
        "--subjects",
        type=str,
        nargs="+",
        required=True,
        help="要创建目录的受试者ID列表 (必需)",
    )
    parser.add_argument(
        "--dry-run", 
        action="store_true", 
        help="模拟运行，只显示将要创建的目录结构，不实际创建"
    )
    parser.add_argument(
        "--verbose", "-v", 
        action="store_true", 
        help="详细输出"
    )
    
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()
    
    # 设置日志级别
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    # 初始化目录创建器
    creator = SubjectDirectoryCreator(args.base_dir)
    
    # 创建目录结构
    if creator.create_subject_directories(args.subjects, args.dry_run):
        if not args.dry_run:
            print("\n🎉 目录创建成功!")
            print("\n现在你可以:")
            print("1. 将预测图像放入 prediction/{subject_id}/ 的相应子目录中")
            print("2. 将分割掩膜放入 segmentation/{subject_id}/ 目录中")
            print("3. 将参考图像放入 reference/{subject_id}/ 目录中")
            print("4. 运行 do_metrics.py 进行图像处理和指标计算")
    else:
        print("\n❌ 目录创建失败!")
        exit(1)


if __name__ == "__main__":
    main()